import argparse
import wandb
from ddos2 import train_and_test_model

sweep_configuration = {
    'method': 'grid',
    'name': 'sweep',
    'metric': {'goal': 'maximize', 'name': 'avg_valid_acc'},
    'parameters':
    {
        'time_steps': {'values': [2, 3, 4, 5, 10, 20, 30, 50, 80, 100, 150]},
        'num_layers': {'values': [1, 2, 3, 4]},
        'hidden_size': {'values': [4, 8, 16, 32, 64, 128, 256]},
        'LSTM_out_size': {'values': [1, 2, 4, 8, 16, 32, 64, 128, 256]},
        'batch_size': {'values': [64, 128, 256, 512, 1024]},
        'epochs': {'min': 3, 'max': 10, 'distribution': 'int_uniform'},
        'lr': {'max': 0.1, 'min': 0.0001},
        'use_cuda': {'values': True},
     }
}

sweep_id = wandb.sweep(
  sweep=sweep_configuration,
  project='ddos-sweep'
  )

# parser = argparse.ArgumentParser()
#
#
# parser.add_argument('--time_steps', type=int, default=10, help='time steps for LSTM')
# parser.add_argument('--num_layers', type=int, default=2, help='the number of layers of LSTM')
# parser.add_argument('--hidden_size', type=int, default=32, help='the hidden size of LSTM')
# parser.add_argument('--LSTM_out_size', type=int, default=2, help='The size of the LSTM output is compressed')
# parser.add_argument('--batch_size', type=int, default=512, help='batch size')
# parser.add_argument('--epochs', type=int, default=10, help='training epochs')
# parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
# parser.add_argument('--use_cuda', type=bool, default=True, help='whether using gpu or cpu')
# parser.add_argument('--random_flag', type=bool, default=False, help='whether using random seed or not')

# args = parser.parse_args()
# def set_random_seed(np_seed, torch_seed):
#     np.random.seed(np_seed)
#     torch.manual_seed(torch_seed)
#     torch.cuda.manual_seed(10)
#     torch.cuda.manual_seed_all(torch_seed)
#
#
# if not args.random_flag:
#     set_random_seed(304, 2021)

# train_and_test_model(args)

def main():
    run = wandb.init()
    config = wandb.config
    train_and_test_model(config)
    run.finish()

# Start sweep job
wandb.agent(sweep_id, function=main)